# app/utils/frontend_manager.py
import asyncio
import logging
from typing import Dict

from fastapi import WebSocket
from starlette.websockets import WebSocketDisconnect

logger = logging.getLogger(__name__)

class FrontendManager:
    def __init__(self):
        self.active_connections: Dict[str, WebSocket] = {}
        self.lock = asyncio.Lock()

    async def connect_frontend(self, websocket: WebSocket):
        """前端连接"""
        await websocket.accept()
        async with self.lock:
            self.active_connections[id(websocket)] = websocket
        logger.info("前端连接")

    async def disconnect_frontend(self, websocket: WebSocket):
        """前端断开连接"""
        async with self.lock:
            if id(websocket) in self.active_connections:
                del self.active_connections[id(websocket)]
        logger.info("前端断开连接")

    async def broadcast(self, message: dict):
        """广播消息给所有前端"""
        async with self.lock:
            for websocket in self.active_connections.values():
                try:
                    await websocket.send_json(message)
                except Exception as e:
                    logger.error(f"广播消息失败: {e}")

# 全局前端管理器
frontend_manager = FrontendManager()

async def handle_frontend_websocket(websocket: WebSocket):
    """处理前端WebSocket连接"""
    await frontend_manager.connect_frontend(websocket)
    try:
        # 当前端连接时，发送所有设备状态
        from app.core.database import get_all_devices
        all_devices = await get_all_devices()
        await websocket.send_json({
            "type": "initial_data",
            "data": all_devices
        })

        # 保持连接
        while True:
            await websocket.receive_text()

    except WebSocketDisconnect:
        logger.info("前端断开连接")
        await frontend_manager.disconnect_frontend(websocket)